from Yolo.darknet import Darknet19
import pickle

YOLOINPUT_FILE = '/mnt/md0/Qiu/ExWorkspace/YOLO_TEST/YOLOINPUT.pkl'
with open(YOLOINPUT_FILE,'rb') as f:
    D = pickle.load(f)

im_data    = D['im_data']
gt_boxes   = D['gt_boxes']
gt_classes = D['gt_classes']
dontcare   = D['dontcare']
size_index = D['size_index']

gt_classes01 = []
for item in gt_classes:
    item[(item > 1).nonzero()] = 0
    gt_classes01.append(item)

model = Darknet19().cuda()
out = model(im_data,gt_boxes,gt_classes01,dontcare,size_index)
#
# print(len(out))
# print(out[0].size())
# print(out[1].size())
# print(out[2].size())
